import argparse
import torch
import os
from mindspore import load_checkpoint
# from models_torch import convert_ms_checkpoint
def convert_ms_checkpoint_notorch(ms_ckpt):
    converted_dict = {}
    for k, v in ms_ckpt.items():
        if not k.startswith('vit'):
            continue
        if 'norm' in k:
            k = k.replace("gamma", "weight")
            k = k.replace("beta", "bias")
        k = k.replace('vit.', '')
        k = k.replace('cls_tokens', 'cls_token')
        k = k.replace('mask_tokens', 'mask_token')
        converted_dict[k] = v.asnumpy()
    return converted_dict

#
parser = argparse.ArgumentParser()

parser.add_argument('ckpt', type=str)
args = parser.parse_args()

if os.path.isdir(args.ckpt):
    jobs = [os.path.join(args.ckpt, f) for f in os.listdir(args.ckpt) if '.ckpt' in f]
else:
    jobs = [args.ckpt,]

for src in jobs:
    dst = src.replace('.ckpt', '.pth')

    ckpt = convert_ms_checkpoint_notorch(load_checkpoint(src))
    print(f'ckpt loaded and converted to np')

    def convert_np_checkpoint_totorch(np_ckpt):
        converted_dict = {}
        for k,v in np_ckpt.items():
            converted_dict[k] = torch.from_numpy(v)
        return converted_dict

    ckpt = convert_np_checkpoint_totorch(ckpt)

    torch.save(ckpt, dst)